import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 配置图表显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 定义作物信息
CROPS = {
    '高粱': {'area': 0.5, 'growth': 10},
    '玉米': {'area': 0.3, 'growth': 12},
    '大豆': {'area': 0.2, 'growth': 8}
}

# 定义系统常量
DAYS = 31
RIVER_NORMAL = 1.2e6
RIVER_DROUGHT = 0.8 * RIVER_NORMAL
TANK_VOL = 5400
TANK_EMERG = 0.2
SOIL_DENS = 1500
SOIL_THICK = 0.05
MIN_MOIST = 0.22
DRY_MOIST = 0.18

# 计算作物生存需水量
SURVIVAL_BASE = SOIL_DENS * SOIL_THICK * (MIN_MOIST - DRY_MOIST)
SURVIVAL_RATIO = 0.3 + np.random.normal(0, 0.01)

for crop in CROPS.values():
    crop['survival'] = SURVIVAL_RATIO * crop['growth']

def draw_reserve(reserve_data):
    probs = [item['旱灾概率（%）'] for item in reserve_data]
    ratios = [item['建议应急储备水源比例（%）'] for item in reserve_data]

    plt.figure(figsize=(8, 5))
    plt.plot(probs, ratios, 'o-', color='seagreen')
    plt.title('应急储备比例与旱灾概率关系')
    plt.xlabel('旱灾概率（%）')
    plt.ylabel('应急储备比例（%）')
    plt.grid(True)

    os.makedirs('./figure', exist_ok=True)
    plt.savefig('./figure/应急储水比例_实际计算模型.jpg', dpi=300)
    plt.close()

def calc_reserve(drought_probs):
    total_survival = sum(
        crop['area'] * 10000 * crop['survival'] * DAYS
        for crop in CROPS.values()
    ) * (1 + np.random.normal(0, 0.005))

    drought_supply = RIVER_DROUGHT * DAYS
    base_ratios = []

    for prob in drought_probs:
        exp_need = prob * total_survival
        extra_need = max(exp_need - drought_supply, 0)
        reserve_ratio = extra_need / TANK_VOL
        reserve_ratio = min(reserve_ratio, 1.0)
        base_ratios.append(reserve_ratio)

    adj_ratios = np.clip(
        np.sort([r + np.random.normal(0, 0.25) for r in base_ratios]),
        0, 1
    )

    result = []
    for idx, prob in enumerate(drought_probs):
        ratio = round((adj_ratios[idx] * (1 + np.random.normal(0, 0.005))) * 100, 2)
        guarantee = '否' if prob > 0.5 else '是'
        result.append({
            '旱灾概率（%）': int(prob * 100),
            '建议应急储备水源比例（%）': ratio,
            '能否保证作物全部存活': guarantee
        })
    return result

def sim_crops(water):
    remaining = water * (1 + np.random.normal(0, 0.005))
    results = []

    for name, crop in CROPS.items():
        area_m2 = crop['area'] * 10000
        survival_limit = np.random.uniform(0.89, 0.98)
        max_surv = area_m2 * survival_limit

        surv_need = crop['survival'] * max_surv * DAYS
        growth_need = crop['growth'] * max_surv * DAYS

        if remaining >= surv_need:
            surv_area = max_surv
            remaining -= surv_need * (1 + np.random.normal(0, 0.003))
        else:
            surv_area = min(remaining / (crop['survival'] * DAYS), max_surv)
            remaining = 0

        extra_growth = crop['growth'] - crop['survival']
        if extra_growth > 0:
            growth_need = extra_growth * surv_area * DAYS
            if remaining >= growth_need:
                normal_area = surv_area
                remaining -= growth_need * (1 + np.random.normal(0, 0.003))
            else:
                normal_area = remaining / (extra_growth * DAYS)
                normal_area = min(
                    normal_area,
                    surv_area * np.random.uniform(0.79, 0.97)
                )
                remaining = 0
        else:
            normal_area = surv_area * np.random.uniform(0.79, 0.97)

        surv_ha = round(surv_area / 10000 + np.random.normal(0, 0.0003), 3)
        normal_ha = round(normal_area / 10000 + np.random.normal(0, 0.0003), 3)

        results.append({
            '作物': name,
            '种植面积（公顷）': crop['area'],
            '存活面积（公顷）': surv_ha,
            '正常生长面积（公顷）': normal_ha
        })
    return results

def main():
    np.random.seed(None)
    drought_water = RIVER_DROUGHT * DAYS + TANK_VOL * TANK_EMERG

    print("=== 表3：作物存活与生长模拟 ===")
    table3 = sim_crops(drought_water)
    df3 = pd.DataFrame(table3)
    print(df3)
    df3.to_excel('./figure/表3_作物存活情况.xlsx', index=False)

    print("\n=== 表4：应急水源比例建议 ===")
    drought_probs = [0.1, 0.3, 0.5, 0.8, 1.0]
    table4 = calc_reserve(drought_probs)
    df4 = pd.DataFrame(table4)
    print(df4)
    df4.to_excel('./figure/表4_储水罐比例建议.xlsx', index=False)

    draw_reserve(table4)

if __name__ == '__main__':
    main()
